# -*- coding: UTF-8 -*-

# FileName     : response
# Author       : EastsunW eastsunw@foxmail.com
# Create at    : 2022-04-22 18:19
# Last Modified: 2022-04-22 18:19
# Modified By  : EastsunW
# -------------
# Description  : 用于TF的自动补全，输入部分名字，返回匹配结果
# -------------

import re
from flask_restful import Resource, fields, marshal_with, reqparse
from Server.Utils.MongoDB import mongo


class Query_FIMO_Class(Resource):

    # 表格的格式
    FIMO_table_fields = {
        "mnvid": fields.String,
        "chr": fields.String,
        "pos": fields.String,
        "strand": fields.String,
        "ref": fields.String,
        "alt": fields.String,
        "rsid": fields.String,
        "transcript": fields.String,
        "gene_symbol": fields.String,
        "tf": fields.String,
        "motif": fields.String,
        "score": fields.Float,
        "pvalue": fields.Float,
        "combinded_start": fields.Integer,
        "combinded_end": fields.Integer,
        "combinded_seq": fields.String,
        "effect": fields.String,
    }
    # 返回数据的格式
    response_fields = {
        "mnv": fields.Integer(0),
        "tf": fields.Integer(0),
        "gene": fields.Integer(0),
        "gain": fields.Integer(0),
        "loss": fields.Integer(0),
        "total": fields.Integer(0),
        "details": fields.List(fields.Nested(FIMO_table_fields))
    }

    def condition_generator(self, argument):
        """根据查询参数构造查询语句
        Args:
            argument (dict): 查询参数
        Returns:
            dict: 查询语句
        """

        condition = {}
        # 精确查询，优先级最高
        if argument["mnv"]:
            bare_id = argument["mnv"].split("MNV")[1]
            if bare_id.startswith("X"):
                mnv_chr = "X"
            else:
                mnv_chr = str(int(bare_id[0:2]))
            condition["$and"] = [
                {"chr": mnv_chr},
                {"mnvid": argument["mnv"]}
            ]
        # 没有给ID，就是模糊查找
        else:
            # range查找
            if argument["range"]:
                parsed_positions = [int(item) for item in re.split(
                    '[chr:-]', argument["range"]) if item]
                assert parsed_positions[0] in list(range(1, 23))
                if len(parsed_positions) == 2:
                    chr, start, end = parsed_positions, None
                elif len(parsed_positions) == 3:
                    chr, start, end = parsed_positions
                else:
                    pass
                condition = {
                    "$and": [
                        {
                            "chr": str(chr)
                        },
                        {
                            "start": {
                                "$lte": end
                            }
                        }
                    ]
                }
                if end:
                    condition["$and"].append({
                        "end": {
                            "$gte": start
                        }
                    })
        # gene查找
        if argument["gene"]:
            condition["gene_symbol"] = argument["gene"]
        # TF查找
        if argument["tf"]:
            condition["tf"] = argument["tf"]
        return condition

    @marshal_with(response_fields)
    def post(self):
        """返回表格数据的接口
        Arguments：
            mnv     (str[])      : 精确查询-MNV
            tf      (str[])      : 模糊查询-转录因子
            gene    (str[])      : 模糊查询-基因
            range   (str)        : 模糊查询-范围
            page    (int)        : 分页-当前页面, 默认1
            size    (int)        : 分页-页面大小, 默认15
            sort    (str)        : 排序依据, 默认不排
            descend (0|1)        : 降序排列(1)
        Returns:
            {
                mnv  : 0,
                tf   : 0,
                gene : 0,
                gain : 0,
                loss : 0,
                total: 0,
                data : Object[]
            }
        """
        # 获取参数：当前页码，每页结果数，查询类型和查询名称
        parser = reqparse.RequestParser()
        parser.add_argument(
            'mnv', type=str,
            required=False,
            nullable=True
        )
        parser.add_argument(
            'tf', type=str,
            required=False,
            nullable=True,
            default=[]
        )
        parser.add_argument(
            'gene', type=str,
            required=False,
            nullable=True,
            default=[]
        )
        parser.add_argument(
            'range', type=str,
            required=False,
            nullable=True
        )
        parser.add_argument(
            'page', type=int,
            required=False,
            nullable=False,
            default=1
        )
        parser.add_argument(
            'size', type=int,
            choices=[15, 50, 100, 0],
            required=False,
            nullable=False,
            default=15
        )
        parser.add_argument(
            'sort', type=str,
            default="pvalue",
            required=False,
            nullable=True,
        )
        parser.add_argument(
            'descend', type=int,
            required=False,
            default=1,
            nullable=False,
        )
        request_args = parser.parse_args()
        assert request_args["mnv"] or \
            request_args["tf"] or \
            request_args["gene"] or \
            request_args["range"]
        condition = self.condition_generator(request_args)
        raw_result = mongo.db.get_collection("fimo").find(condition)
        if request_args["sort"]:
            raw_result = raw_result.sort(
                [(request_args["sort"], 1 if request_args["descend"] == 0 else -1)])
        if request_args["size"]:
            num_skip = (request_args["page"] - 1) * request_args["size"]
            response_result = list(raw_result.skip(
                num_skip).limit(request_args["size"]))
        else:
            response_result = list(raw_result)

        # 获取统计信息
        num_total = raw_result.count()
        num_mnv = len(raw_result.distinct("mnvid"))
        num_tf = len(raw_result.distinct("tf"))
        num_gene = len(raw_result.distinct("gene_symbol"))
        num_type_raw = list(mongo.db.get_collection("fimo").aggregate(
            [
                {
                    "$match": condition
                },
                {
                    "$group": {
                        "_id": "$effect",
                        "count": {
                            "$sum": 1
                        }
                    }
                },
                {
                    "$group": {
                        "_id": "null",
                        "counts": {
                            "$push": {
                                "k": "$_id",
                                "v": "$count"
                            }
                        }
                    }
                },
                {
                    "$replaceRoot": {
                        "newRoot": {
                            "$arrayToObject": "$counts"
                        }
                    }
                }
            ]))
        if len(num_type_raw) >= 1:
            num_type = num_type_raw[0]
            return {
                "mnv": num_mnv,
                "tf": num_tf,
                "gene": num_gene,
                "gain": num_type["Gain"] if "Gain" in num_type else 0,
                "loss": num_type["Loss"] if "Loss" in num_type else 0,
                "total": num_total,
                "details": response_result
            }
        return {
            "mnv": 0,
            "tf": 0,
            "gene": 0,
            "gain": 0,
            "loss": 0,
            "total": 0,
            "details": []
        }
